
"""
电子小说章节解析
"""

import re
from typing import List
from models import Chapter

volume_reg = re.compile("^\s*(第[\d零一二三四五六七八九十百千万]+[篇卷])")
chapter_reg = re.compile("^\s*(第[\d零一二三四五六七八九十百千万]+章)")


def parse_novel(novel_texts: List[str]) -> List[Chapter]:
    """小说解析"""
    chapters = []
    volume = ''
    chapter_number = 0
    chapter_title = ''

    contents = []
    for line in novel_texts:
        line = line.strip().replace("\u3000", "")
        if line == '' or line == chapter_title:
            continue
        if volume_reg.search(line) is not None:
            volume = line
            continue
        elif chapter_reg.search(line) is None:
            contents.append(line)
            continue
        chapter = Chapter(
            volume=volume,
            chapter_number=chapter_number,
            chapter_title=chapter_title,
            chapter_content='\n'.join(contents),
        )
        chapters.append(chapter)
        contents = []
        chapter_number += 1
        chapter_title = line
    return chapters


def merge_chapters(chapters: List[Chapter], max_len=64000) -> List[Chapter]:
    """合并章节内容"""
    merged_chapters = []
    volume = ''
    word_count = 0
    batch_contents = []
    batch_titles = []
    for idx, chapter in enumerate(chapters):
        words = f"**[{chapter.chapter_title}]**\n{chapter.chapter_content}"
        if (idx != 0 and chapter.volume != volume) or word_count + len(words) > max_len:
            merged_chapter = Chapter(
                volume=volume,
                chapter_number=len(merged_chapters) + 1,
                chapter_title=merge_titles(batch_titles),
                chapter_content='\n\n'.join(batch_contents),
            )
            merged_chapters.append(merged_chapter)
            word_count = 0
            batch_contents = []
            batch_titles = []
        volume = chapter.volume
        batch_contents.append(words)
        if chapter.chapter_title != '':
            batch_titles.append(chapter.chapter_title)
        word_count += len(words)

    if len(batch_contents) != 0:
        merged_chapter = Chapter(
            volume=volume,
            chapter_number=len(merged_chapters) + 1,
            chapter_title=merge_titles(batch_titles),
            chapter_content='\n'.join(batch_contents),
        )
        merged_chapters.append(merged_chapter)
    return merged_chapters


def merge_titles(batch_titles: List[str]) -> str:
    """合并标题"""
    merged_titles = []
    if len(batch_titles) >= 1:
        found = chapter_reg.search(batch_titles[0])
        if found is not None:
            merged_titles.append(found.group(1))
    if len(batch_titles) > 1:
        found = chapter_reg.search(batch_titles[-1])
        if found is not None:
            merged_titles.append(found.group(1))
    if len(merged_titles) == 0:
        merged_title = f'第{len(merged_chapters) + 1}部分'
    else:
        merged_title = '至'.join(merged_titles)
    return merged_title


if __name__ == '__main__':
    from core.utils.path import data_dir
    file = data_dir / '凡人修仙传(1-500章).txt'
    try:
        novel_texts = open(file, 'r', encoding='GB18030').readlines()
    except:
        novel_texts = open(file, 'r', encoding='utf-8').readlines()
    chapters = parse_novel(novel_texts)
    merged_chapters = merge_chapters(chapters[0:100], max_len=32000)
    for idx, chapter in enumerate(merged_chapters):
        print(idx, ' --- ',chapter.chapter_title, ' --- ', [chapter.chapter_content])